#include "lbm2d.h"

Cell::Cell(void) {
	for (int l = 0; l<Q_DIM; l++) {
		f[l] = 0.0;
		fstream[l] = 0.0;
		feq[l] = 0.0;
	}
}

double Cell::rho(void) { //rho calculado con f (densidad de poblacion total de la celda)
	double ret = 0.0;
	for(int l=0; l<Q_DIM; l++)
		ret = ret + fstream[l];
	return ret;
// 	return 1.0;
}

double Cell::rhoOld(void) { //rho calculado con f (densidad de poblacion total de la celda)
	double ret = 0.0;
	for(int l=0; l<Q_DIM; l++)
		ret = ret + f[l];
	return ret;
// 	return 1.0;
}

vec<double> Cell::u(void) { //u, velocidad media de la celda
	double _rho = rho();
	vec<double> ret = {0, 0};
	if (_rho > 0) {
		for(int l=0; l<Q_DIM; l++) {
			ret.x = ret.x + eVector[l].x*fstream[l];
			ret.y = ret.y + eVector[l].y*fstream[l];
		}
		ret.x = ret.x / _rho;
		ret.y = ret.y / _rho;
	}
	return ret;
}

vec<double> Cell::uOld(void) {
	double _rho = rhoOld();
	vec<double> ret = {0, 0};
	if (_rho > 0) {
		for(int l=0; l<Q_DIM; l++) {
			ret.x = ret.x + eVector[l].x*f[l];
			ret.y = ret.y + eVector[l].y*f[l];
		}
		ret.x = ret.x / _rho;
		ret.y = ret.y / _rho;
	}
	return ret;
}

double Cell::getFeq(double rho, vec<double> u, int l) {
	double a = eVector[l]*u;
	return wCoeff[l]*rho*(1.0 - 1.5*pow(u.norm(), 2) + 3.0*a + 4.5*pow(a,2));
}

void Cell::setU(vec<double> u, double rho) {
	for(int l=0; l<Q_DIM; l++)
		f[l] = getFeq(rho, u, l);
}

ostream &operator<<(ostream &os, const Cell& cell) {
	double ret = 0.0;
	for(int l=0; l < Q_DIM; l++)
		ret = ret + cell.f[l];
	os << ret;
	return os;
}

LBM::LBM() {
}

LBM::LBM(vec<int> dim, double omega) : dim(dim), omega(omega) {
	grid = vector< vector<Cell> >(dim.x, vector<Cell>(dim.y));
	tau = 1.0/omega;
	viscosity = (tau - 0.5) / 3.0;
	periodicBoundary = {true, true};
}

void LBM::initialize(double rho) {
	vec<double> u0 = {0.2, 0.0};
	vec<double> u1 = {-0.2, 0.0};
	vec<double> u_null = {0.0, 0.0};

	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++) {
				grid[i][j].isObstacle = false;
// 				grid[i][j].setU(u0, rho);
// 				if ((j == 0) || (j == dim.y-1)) grid[i][j].isObstacle = true;
				if (j <= dim.y/2) grid[i][j].setU(u0, rho);
					else grid[i][j].setU(u1, rho);
		}
}

double LBM::getUMax() {
	double uMax = 0;
	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++) {
			vec<double> _u = grid[i][j].uOld();
			double uMag = sqrt(pow(_u.x,2) + pow(_u.y,2));
			if (uMag > uMax)
				uMax = uMag;
		}
	if (uMax <= 0) uMax = 1;
	return uMax;
}

void LBM::putRectangleObstacle(int x, int y, int dx, int dy) {
	for(int i=x; i<x+dx; i++)
		for(int j=y; j<y+dy; j++)
			grid[i][j].isObstacle = true;
}

void LBM::applyForce() {
	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++) {
			vec<double> _u = grid[i][j].u();
// 			double f = 8.0 * viscosity * _u.x * grid[i][j].rho() / (6.0 * pow(dim.y,2));
			double f = 8.0 * viscosity * 0.1 * 1.0 / (6.0 * pow(dim.y,2));
			grid[i][j].f[East] += f;
			grid[i][j].f[NorthEast] += f;
			grid[i][j].f[SouthEast] += f;

			grid[i][j].f[West] -= f;
			grid[i][j].f[NorthWest] -= f;
			grid[i][j].f[SouthWest] -= f;
		}
}


void LBM::stream_step() {
	for (int i = 0; i < dim.x; i++) {
		for (int j = 0; j < dim.y; j++) {
			for (int l = 0; l < Q_DIM; l++) {
				int lInv = invertedIndices[l];
				vec<int> streamCell = {i + (int)eVector[lInv].x, j + (int)eVector[lInv].y};

				//condiciones periódicas en x
				if (periodicBoundary.x == true) {
					if (streamCell.x < 0)
						streamCell.x = dim.x-1;
					else if (streamCell.x == dim.x)
						streamCell.x = 0;
				}
				//condiciones periódicas en y
				if (periodicBoundary.y == true) {
					if (streamCell.y < 0)
						streamCell.y = dim.y-1;
					else if (streamCell.y == dim.y)
						streamCell.y = 0;
				}

				grid[i][j].fstream[l] = 0.0;
				if (grid[streamCell.x][streamCell.y].isObstacle)
					grid[i][j].fstream[l] = grid[i][j].f[lInv];
				else
					grid[i][j].fstream[l] = grid[streamCell.x][streamCell.y].f[l];
			}
		}
	}
}

void LBM::collide_step() {
	for (int i = 0; i < dim.x; i++) {
		for (int j = 0; j < dim.y; j++) {
			double _rho = grid[i][j].rho();
			vec<double> _u = grid[i][j].u();
			for (int l = 0; l < Q_DIM; l++) {
				grid[i][j].feq[l] = grid[i][j].getFeq(_rho, _u, l);
				grid[i][j].f[l] = (1.0-omega)*grid[i][j].fstream[l] + omega*grid[i][j].feq[l];
			}
		}
	}
}

double LBM::norm() {
	double ret = 0.0;
	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++)
			for (int l = 0; l < Q_DIM; l++)
				if (grid[i][j].isObstacle == false) ret += grid[i][j].f[l];
	return ret;
}

ostream &operator<<(ostream &os, const LBM& lbm) {
	for (int j = 0; j < lbm.dim.y; j++) {
		os << endl;
		for (int i = 0; i < lbm.dim.x; i++) os << lbm.grid[i][j] << " ";
	}
	return os;
}
